#include <cuda_runtime.h>
#include <stdio.h>
#include <stdlib.h>

// float4的作用：进行向量化访问的优点
// float4是cuda的内置类型，通过一次内存事务读写4个float数据，可以减少访存次数
#define FLOAT(a) *(float4*)(&(a))

// 上取整
#define CEIL(a, b) ((a + b - 1)/(b))

// 报错日志，预处理宏不需要作为参数传递
#define cudaCheck(err) _cudaCheck(err, __FILE__, __LINE__)

inline void _cudaCheck(cudaError_t error, const char *file, int line) {
    if(error != cudaSuccess) {
        // cudaGetErrorString() 
        printf("[CUDA ERROR] at file %s(line %d):\n%s\n", file, line, cudaGetErrorString(error));
        exit(EXIT_FAILURE);
    }
    return ;
}

__global__ void elementwise_add_float4(float *a, float *b, float *c, int N) {
    int idx = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
    // 超出了最大数据量线程不需要工作
    if (idx >= N) {
        return ;
    }

    float4 temp_a = FLOAT(a[idx]);
    float4 temp_b = FLOAT(b[idx]);
    float4 temp_c;
    temp_c.x = temp_a.x + temp_b.x;
    temp_c.y = temp_a.y + temp_b.y;
    temp_c.z = temp_a.z + temp_b.z;
    temp_c.w = temp_a.w + temp_b.w;

    FLOAT(c[idx]) = temp_c;
}

int main()
{
    constexpr int N = 7;
    float *a_h = (float *)malloc(N * sizeof(float));
    float *b_h = (float *)malloc(N * sizeof(float));
    float *c_h = (float *)malloc(N * sizeof(float));
    for (int i = 0; i < N; ++i) {
        a_h[i] = i;
        b_h[i] = N - i - 1;
    }

    float *a_d = nullptr;
    float *b_d = nullptr;
    float *c_d = nullptr;

    cudaCheck(cudaMalloc((void**)&a_d, N * sizeof(float)));
    cudaCheck(cudaMalloc((void**)&b_d, N * sizeof(float)));
    cudaCheck(cudaMalloc((void**)&c_d, N * sizeof(float)));
    cudaCheck(cudaMemcpy(a_d, a_h, N * sizeof(float), cudaMemcpyHostToDevice));
    cudaCheck(cudaMemcpy(b_d, b_h, N * sizeof(float), cudaMemcpyHostToDevice));

    int block_size = 256;
    int grid_size = CEIL(CEIL(N, 4), 1024);
    elementwise_add_float4<<<grid_size, block_size>>>(a_d, b_d, c_d, N);

    cudaCheck(cudaMemcpy(c_h, c_d, N * sizeof(float), cudaMemcpyDeviceToHost));
    printf("a_h:\n");
    for (int i = 0; i < N; i++ ) {
        if (i == N-1) printf("%f\n", a_h[i]);
        else printf("%f ", a_h[i]);
    }
    printf("b_h:\n");
    for (int i = 0; i < N; i++ ) {
        if (i == N-1) printf("%f\n", b_h[i]);
        else printf("%f ", b_h[i]);
    }
    printf("c_h:\n");
    for (int i = 0; i < N; i++ ) {
        if (i == N-1) printf("%f\n", c_h[i]);
        else printf("%f ", c_h[i]);
    }

    free(a_h); free(b_h); free(c_h);
    cudaCheck(cudaFree(a_d));
    cudaCheck(cudaFree(b_d));
    cudaCheck(cudaFree(c_d));
    
    return 0;
}